from __future__ import absolute_import, division, print_function, unicode_literals

import glob
import os
import argparse
import json
import torch
from scipy.io.wavfile import write
from hifi.env import AttrDict
from hifi.meldataset import mel_spectrogram, MAX_WAV_VALUE, load_wav
from hifi.models import Generator

h = None
device = None


def load_checkpoint(filepath, device):
    assert os.path.isfile(filepath)
    print("Loading '{}'".format(filepath))
    checkpoint_dict = torch.load(filepath, map_location=device)
    print("Complete.")
    return checkpoint_dict


def get_mel(x):
    return mel_spectrogram(x, h.n_fft, h.num_mels, h.sampling_rate, h.hop_size, h.win_size, h.fmin, h.fmax)


def scan_checkpoint(cp_dir, prefix):
    pattern = os.path.join(cp_dir, prefix + '*')
    cp_list = glob.glob(pattern)
    if len(cp_list) == 0:
        return ''
    return sorted(cp_list)[-1]


def inference(x):
    generator = Generator(h).to(device)

    state_dict_g = load_checkpoint(a.checkpoint_file, device)
    generator.load_state_dict(state_dict_g['generator'])


    generator.eval()
    generator.remove_weight_norm()
    y_g_hat = generator(x)
    audio = y_g_hat.squeeze()
    print("audio",audio.shape)
    audio = audio * MAX_WAV_VALUE
    audio = audio.cpu().numpy().astype('int16')

           


def Init():
    print('Initializing Inference Process..')

    # parser = argparse.ArgumentParser()
    # # parser.add_argument('--input_wavs_dir', default='test_files')
    # # parser.add_argument('--output_dir', default='generated_files')
    # parser.add_argument('--checkpoint_file', required=True)
    # a = parser.parse_args()
    
    checkpoint_file="/opt/data/private/VC/hifi-gan/checkpoints/VCTK_V1/generator_v1"

    config_file = os.path.join(os.path.split(checkpoint_file)[0], 'config.json')
    with open(config_file) as f:
        data = f.read()

    global h
    json_config = json.loads(data)
    h = AttrDict(json_config)

    generator = Generator(h).to(device)

    state_dict_g = load_checkpoint(checkpoint_file, device)
    generator.load_state_dict(state_dict_g['generator'])


    generator.eval()
    generator.remove_weight_norm()
    
    
    # x=torch.randn(3,80,100)
    # print(x)
    # inference(x)
    
    # print(generator)
    return generator
    # torch.manual_seed(h.seed)
    # global device
    # if torch.cuda.is_available():
    #     torch.cuda.manual_seed(h.seed)
    #     device = torch.device('cuda')
    # else:
    #     device = torch.device('cpu')

    # x=torch.randn(3,80,100)
    # inference(x)

Init()


